#include "Algorithm.hpp"
#include <iostream>
#include <cstring>
#include <math.h>
#include "utils.h"
#include "parallel.hpp"
using namespace std;

Number Algorithm::toleranceLevel = 1e-9;
Index Algorithm::maxSteps = 500;
bool Algorithm::_checkInternalDerivatives = false;
unsigned int Algorithm::verboseLevel = 0;
bool Algorithm::solveWithIpopt = false;

Algorithm::Algorithm(ProblemRef problemDefinition, Index subProblemsCount) :
	problemDefinition(problemDefinition), subProblemsCount(subProblemsCount),
			dim(-1), restrictions(-1), timeParallel(0), timeParallelSummed(0),
			timeOverall(0), timeSynchro(0) {
	Index tmp;
	TNLP::IndexStyleEnum indexingStyle;

	problemDefinition.get_nlp_info(dim, restrictions, orgNnzJac, tmp,
			indexingStyle);
	assert(indexingStyle == TNLP::C_STYLE);

	allocateStructures();

	//punkt startowy
	problemDefinition.get_starting_point(dim, true, lastOptX, false, NULL,
			NULL, restrictions, false, NULL);

	//zapamietaj oryginalne ograniczenia
	problemDefinition.get_bounds_info(dim, x_l, x_u, restrictions, g_l, g_u);

	//compute indexes for procs
	problemDefinition.get_blocks_indexes(dim, restrictions, subProblemsCount,
			procStartIndex);
	procStartIndex[subProblemsCount] = dim;

	//oblicz parametry jakobianu
	computeJacobianData();

	//inicjalizacja podproblemów równoległych i synchronizacji
	try {
		for (int p = 0; p < subProblemsCount; ++p) {
			subProblems[p] = new SubProblem(*this, p, procStartIndex[p],
					procStartIndex[p + 1]);
			subProblems[p]->AddRef(NULL);
			appInit(&apps[p], false);
			apps[p].AddRef(NULL);
		}
		appInit(&apps[subProblemsCount], true);
		apps[subProblemsCount].AddRef(NULL);
	} catch (const bad_alloc& x) {
		deallocateStructures();
		throw ;
	}

	step = 0;
}

void Algorithm::computeJacobianData() {
	//musi byc wywolane po policzeniu indeksow dla procesorow

	//pobierz oryginalne polozenia
	problemDefinition.eval_jac_g(dim, NULL, false, restrictions, orgNnzJac,
			jacIRow, jacJCol, NULL);

	//wypelnij tablice ilosci niezerowych elementow w poszczegolnych blokach
	memset(procToJacElemNum, 0, subProblemsCount * sizeof(Index));
	memset(procToJacBlockNum, 0, subProblemsCount * sizeof(Index));
	memset(restrToProcJacBlockInd, -1, restrictions * sizeof(Index));

	for (Index nzi = 0; nzi < orgNnzJac; ++nzi) {
		Index proc = varToProc(jacJCol[nzi]);
		procToJacElemNum[proc]++;
		if (restrToProcJacBlockInd[jacIRow[nzi]] == -1) { //bo kazde ograniczenie w tylko jednym bloku
			restrToProcJacBlockInd[jacIRow[nzi]] = procToJacBlockNum[proc]++;
		}
	}

	//dystrybuanta procToJacBlockNum
	procToJacBlockNumDist[0] = 0;
	for (Index proc = 1; proc <= subProblemsCount; proc++) {
		procToJacBlockNumDist[proc] = procToJacBlockNumDist[proc - 1]
				+ procToJacBlockNum[proc - 1];
	}
}

Index Algorithm::varToProc(Index varInd) const {
	for (Index proc = 0; proc < subProblemsCount; ++proc) {
		if (varInd < procStartIndex[proc + 1]) {
			return proc;
		}
	}
	assert(false);
}

Index Algorithm::nnzJacIndToProc(Index orgNzi) const {
	return varToProc(jacJCol[orgNzi]);
}

void Algorithm::appInit(IpoptApplication *app, bool synchro) {
	if (synchro) {
		app->Options()->SetNumericValue("tol", Algorithm::toleranceLevel
				* Algorithm::toleranceLevel);
	} else {
		app->Options()->SetNumericValue("tol", Algorithm::toleranceLevel * 0.1);
	}
	app->Options()->SetStringValue("hessian_approximation", "limited-memory");
	app->Options()->SetIntegerValue("max_iter", 1000);
	if (verboseLevel > 0)
		app->Options()->SetIntegerValue("print_level", verboseLevel - 1);
	else
		app->Options()->SetIntegerValue("print_level", 0);

	if (Algorithm::_checkInternalDerivatives)
		app->Options()->SetStringValue("derivative_test", "first-order");

	app->Initialize();
}

IpoptApplication* Algorithm::appFactory() {
	IpoptApplication* app = new IpoptApplication();
	appInit(app, false);

	return app;
}

void Algorithm::checkDerivatives() {
	cout << "Derivative check!\n";

	SmartPtr<IpoptApplication> app = new IpoptApplication();
	app->Options()->SetNumericValue("tol", Algorithm::toleranceLevel);
	app->Options()->SetStringValue("hessian_approximation", "limited-memory");
	app->Options()->SetStringValue("derivative_test", "first-order");
	app->Options()->SetIntegerValue("max_iter", 0);
	app->Initialize();

	problemDefinition.AddRef(NULL);
	app->OptimizeTNLP(&problemDefinition);
	cout << "End of derivative check!\n";
}

Algorithm::~Algorithm() {
	deallocateStructures();
}

ApplicationReturnStatus Algorithm::solve() {
	long ts = getTimeMilis();

	//normalized gradient as optimality function
	computeOptimalityFunction();

	do {
		doStep();
	} while (optimalityFun > toleranceLevel && step < maxSteps);

	timeOverall = getTimeMilis() - ts;
	timeSynchro = timeOverall - timeParallel;

	if (step == maxSteps)
		return Maximum_Iterations_Exceeded;

	return Solve_Succeeded;
}

#define MIN_MAX 1e9

void Algorithm::computeOptimalityFunction() {
	//normalized gradient as optimality function
	problemDefinition.eval_grad_f(dim, lastOptX, true, optFunDiagMatrix);

	gradientNorm = 0;
	for (int i = 0; i < dim; ++i) {
		gradientNorm += optFunDiagMatrix[i] * optFunDiagMatrix[i];
	}

	gradientNorm = sqrt(gradientNorm);

	Number grDiv = 0;
	if (gradientNorm != 0) {
		grDiv = -1.0 / gradientNorm;
		assert(finite(grDiv));
	}
	for (int i = 0; i < dim; ++i) {
		optFunDiagMatrix[i] *= grDiv;
	}

	if (restrictions == 0) {
		optimalityFun = gradientNorm;
	} else {
		//naiwne przeszukiwanie warunku 3.2 (Ferris strona 6)
		//po poszczegolych zmiennych
		Number min = MIN_MAX;
		Number alpha = toleranceLevel; //TODO mozna zmienic ew. na wieksze
		memcpy(xOrgForOpt, lastOptX, sizeof(Number) * dim);

		do {
			for (Index varNum = 0; varNum < dim; varNum++) {
				Number tmp = xOrgForOpt[varNum];
				xOrgForOpt[varNum] += alpha;
				problemDefinition.eval_g(dim, xOrgForOpt, true, restrictions,
						restrForOpt);
				xOrgForOpt[varNum] = tmp;

				bool restrOk = true;
				for (Index restrNum = 0; restrNum < restrictions; restrNum++) {
					if (restrForOpt[restrNum] < g_l[restrNum]
							|| restrForOpt[restrNum] > g_u[restrNum]) {
						restrOk = false;
						break;
					}
				}
				if (restrOk && optFunDiagMatrix[varNum] < min) {
					min = optFunDiagMatrix[varNum];
				}
			}
			alpha = -alpha;
		} while (alpha < 0); //dwa powtorzenia dla alha i -alpha

		if (min == MIN_MAX) {
			//wpadlismy pod ograniczenia, w zadna strone nie mozna sie ruszyc
			optimalityFun = 0; //TODO czy na pewno?
		} else {
			optimalityFun = -min;
		}
	}
}

//podzial i synchronizacja
void Algorithm::doStep() {
	if (Algorithm::verboseLevel) {
		cout << " *** Step: " << (step + 1) << endl;
	}

	//parallely solve subproblems
	if (Algorithm::verboseLevel) {
		cout << " *** Solving subproblems parallely..." << endl;
	}

	int p = 0;
	long _timeParallel = 0;
	long ts = getTimeMilis();

#ifdef WITH_MPI
	// ******************** MPI ****************************
	int rank, size;
	MC(MPI_Comm_rank(MPI_COMM_WORLD, &rank));
	MC(MPI_Comm_size(MPI_COMM_WORLD, &size));

	for (p = rank; p < subProblemsCount; p += size) {
		long _ts = getTimeMilis();
		apps[p].OptimizeTNLP(subProblems[p]);

		_timeParallel += getTimeMilis() - _ts;
	}
	if (rank != 0) {
		//nie root-y wysylaja rozwiazania policzonych podproblemow
		for (p = rank; p < subProblemsCount; p += size) {
			MC(MPI_Ssend(subProblems[p]->xOrgSpace, dim, MPI_DOUBLE,
							0, p, MPI_COMM_WORLD));
		}
	} else {
		//wykonuje tylko root

		for (p = 1; p < subProblemsCount; ++p) {
			if (p % size != 0) { //nie moj podprobem
				MC(MPI_Recv ( subProblems[p]->xOrgSpace, dim, MPI_DOUBLE,
								MPI_ANY_SOURCE, p, MPI_COMM_WORLD, NULL));
			}
		}

		timeParallel += getTimeMilis() - ts;
		timeParallelSummed += _timeParallel;

		//synchronization step
		if (Algorithm::verboseLevel) {
			cout << " *** Synchronization..." << endl;
		}

		apps[subProblemsCount].OptimizeTNLP(synch);

		//read solution
		for (int i = 0; i < dim; ++i) {
			lastOptX[i] = synch->xOrgSpace[i];
		}
	}
	MC(MPI_Bcast(lastOptX, dim, MPI_DOUBLE, 0, MPI_COMM_WORLD));
#else
	// ***************** NOT NOT NOT MPI **************************
#pragma omp parallel for reduction(+:_timeParallel)
	for (p = 0; p < subProblemsCount; ++p) {
		long _ts = getTimeMilis();
		apps[p].OptimizeTNLP(subProblems[p]);

		_timeParallel += getTimeMilis() - _ts;
	}

	timeParallel += getTimeMilis() - ts;
	timeParallelSummed += _timeParallel;

	//synchronization step
	if (Algorithm::verboseLevel) {
		cout << " *** Synchronization..." << endl;
	}

	apps[subProblemsCount].OptimizeTNLP(synch);

	//read solution
	for (int i = 0; i < dim; ++i) {
		lastOptX[i] = synch->xOrgSpace[i];
	}
	// ****************** END NOT MPI ********************
#endif

	//normalized gradient as optimality function
	computeOptimalityFunction();

	if (verboseLevel) {
		cout << "Optimality fun: " << optimalityFun << endl;
	}

	++step;
}

void Algorithm::allocateStructures() {
	assert(dim > 0);
	assert(dim >= 0);

	try {
		optFunDiagMatrix = new Number[dim * 3 + restrictions];
		lastOptX = optFunDiagMatrix + dim;
		xOrgForOpt = lastOptX + dim;
		restrForOpt = xOrgForOpt + dim;

		procStartIndex = new Index[subProblemsCount + 1];

		subProblems = new SubProblemPt[subProblemsCount];

		x_l = new Number[2 * dim + 2 * restrictions];
		x_u = x_l + dim;
		g_l = x_u + dim;
		g_u = g_l + restrictions;

		jacIRow = new Index[orgNnzJac * 2 + 3 * subProblemsCount + 1
				+ restrictions];
		jacJCol = jacIRow + orgNnzJac;
		procToJacElemNum = jacJCol + orgNnzJac;
		procToJacBlockNum = procToJacElemNum + subProblemsCount;
		procToJacBlockNumDist = procToJacBlockNum + subProblemsCount;
		restrToProcJacBlockInd = procToJacBlockNumDist + subProblemsCount + 1;

		synch = new Synchro(*this);
		synch->AddRef(NULL);

		apps = new IpoptApplication[subProblemsCount + 1];

	} catch (const bad_alloc& x) {
		deallocateStructures();
		throw ;
	}
}

void Algorithm::deallocateStructures() {
	if (optFunDiagMatrix != NULL) {
		delete[] optFunDiagMatrix;
		optFunDiagMatrix = NULL;
		lastOptX = NULL;
		xOrgForOpt = NULL;
		restrForOpt = NULL;
	}

	if (procStartIndex != NULL) {
		delete[] procStartIndex;
		procStartIndex = NULL;
	}

	if (x_l != NULL) {
		delete[] x_l;
		x_l = x_u = g_l = g_u = NULL;
	}

	if (jacIRow != NULL) {
		delete[] jacIRow;
		jacIRow = NULL;
		jacJCol = NULL;
		procToJacElemNum = NULL;
		procToJacBlockNum = NULL;
		procToJacBlockNumDist = NULL;
		restrToProcJacBlockInd = NULL;
	}

	if (apps != NULL) {
		delete[] apps;
		apps = NULL;
	}

	if (subProblems != NULL) {
		for (int p = 0; p < subProblemsCount; ++p) {
			if (subProblems[p] != NULL) {
				delete subProblems[p];
				subProblems[p] = NULL;
			}
		}
		delete[] subProblems;
		subProblems = NULL;
	}

	if (synch != NULL) {
		delete synch;
		synch = NULL;
	}
}
